import numpy as np
import matplotlib.pyplot as plt
from sklearn.neural_network import MLPClassifier
from sklearn.preprocessing import StandardScaler
from sklearn.model_selection import train_test_split

data = np.loadtxt(r'../data/ex2data2.txt', delimiter=',')
x = data[:, :-1]
y = data[:, -1]

scaler = StandardScaler()
x = scaler.fit_transform(x)

x_train, x_test, y_train, y_test = train_test_split(x, y, train_size=0.7, random_state=1)

clf = MLPClassifier([100, 50], 'relu', alpha=0.01, max_iter=2000)
clf.fit(x_train, y_train)
print(f'training score = {clf.score(x_train, y_train)}')
print(f'testing score = {clf.score(x_test, y_test)}')

pos_idx = y_train == 1
neg_idx = np.invert(pos_idx)
plt.scatter(x_train[pos_idx, 0], x_train[pos_idx, 1], s=1, c='r', zorder=50, label='pos train')
plt.scatter(x_train[neg_idx, 0], x_train[neg_idx, 1], s=1, c='b', zorder=50, label='neg train')
pos_idx = y_test == 1
neg_idx = np.invert(pos_idx)
plt.scatter(x_test[pos_idx, 0], x_test[pos_idx, 1], s=10, c='r', zorder=10, label='pos test')
plt.scatter(x_test[neg_idx, 0], x_test[neg_idx, 1], s=10, c='b', zorder=10, label='neg test')

x1_min, x1_max = np.min(x[:, 0]), np.max(x[:, 0])
x2_min, x2_max = np.min(x[:, 1]), np.max(x[:, 1])
xx, yy = np.mgrid[x1_min:x1_max:0.01, x2_min:x2_max:0.01]
hh = clf.predict(np.c_[xx.ravel(), yy.ravel()])
hh.resize(xx.shape)
plt.contour(xx, yy, hh > 0.5, zorder=0, cmap=plt.cm.Paired)

plt.grid()
plt.legend()
plt.show()
